Param(
  [Parameter(Mandatory = $true)] [string]$Input,     # 源视频（.mp4/.mov/.avi/...）
  [int]$Seconds = 15,                                 # 截取秒数
  [int]$Fps = 3,                                      # 抽帧 FPS
  [int]$ScaleH = 720,                                 # 抽帧高度（等比缩放）
  [string]$Device = 'cuda',                           # 训练设备：cuda/mps/cpu
  [string]$Weights = 'weights/yolov8n.pt',            # 预训练权重
  [string]$DatasetRoot = 'datasets/yolo',             # 数据集根目录
  [switch]$Clean                                      # 清理历史同名输出
)

$ErrorActionPreference = 'Stop'

function Require-Tool($name) {
  $exists = (Get-Command $name -ErrorAction SilentlyContinue) -ne $null
  return $exists
}

function Invoke-Yolo {
  param([string]$Args)
  if (Require-Tool 'yolo') {
    Write-Host "[YOLO] yolo $Args" -ForegroundColor DarkCyan
    & yolo @Args.Split(' ') | Out-Host
  } else {
    if (-not (Require-Tool 'python')) { throw 'python 不在 PATH，且找不到 yolo 命令。请先激活你的 Python 环境。' }
    Write-Host "[YOLO] python -m ultralytics $Args" -ForegroundColor DarkCyan
    & python -m ultralytics @Args.Split(' ') | Out-Host
  }
}

if (-not (Test-Path $Input)) { throw "找不到输入视频：$Input" }
if (-not (Require-Tool 'ffmpeg')) { throw '未发现 ffmpeg，请先安装：winget install -e --id Gyan.FFmpeg --source winget' }
if (-not (Require-Tool 'python')) { Write-Warning '未发现 python，后续 remap/fix_label 脚本将失败。请先激活 Python 环境。' }

# 路径准备
$tmpDir = Join-Path 'tmp' 'pipeline_from_video'
New-Item -ItemType Directory -Force -Path $tmpDir | Out-Null

$stem = [System.IO.Path]::GetFileNameWithoutExtension($Input)
$clip = Join-Path $tmpDir ("{0}_0_{1}s.mp4" -f $stem, $Seconds)

Write-Host "[1/6] 裁剪视频 → $Seconds s: $clip" -ForegroundColor Cyan
$ff = @('-y','-ss','0','-t',"$Seconds",'-i',$Input,'-c:v','libx264','-an','-preset','veryfast','-crf','23',$clip)
$p = Start-Process -FilePath ffmpeg -ArgumentList $ff -NoNewWindow -PassThru -Wait
if ($p.ExitCode -ne 0) { throw "ffmpeg 裁剪失败：$Input" }

Write-Host "[2/6] 抽帧 + 生成 meta.csv" -ForegroundColor Cyan
$framesRoot = Join-Path $tmpDir 'frames_clip'
if ($Clean -and (Test-Path $framesRoot)) { Remove-Item -Recurse -Force $framesRoot }
New-Item -ItemType Directory -Force -Path $framesRoot | Out-Null
& powershell -ExecutionPolicy Bypass -File (Join-Path 'scripts' 'extract_frames.ps1') `
  -InPath $clip -OutRoot $framesRoot -Fps $Fps -ScaleH $ScaleH -Owner $env:USERNAME

# 展平到 YOLO 数据集结构
$imgTrain = Join-Path $DatasetRoot (Join-Path 'images' 'train')
$imgVal   = Join-Path $DatasetRoot (Join-Path 'images' 'val')
$labTrain = Join-Path $DatasetRoot (Join-Path 'labels' 'train')
$labVal   = Join-Path $DatasetRoot (Join-Path 'labels' 'val')
foreach ($d in @($imgTrain,$imgVal,$labTrain,$labVal)) { New-Item -ItemType Directory -Force -Path $d | Out-Null }

$sub = Get-ChildItem -Directory -Path $framesRoot | Select-Object -First 1
if (-not $sub) { throw "抽帧输出为空：$framesRoot" }

Write-Host "[3/6] 拷贝图片至 YOLO 数据集 (train/val)" -ForegroundColor Cyan
Copy-Item -Path (Join-Path $sub.FullName '*.jpg') -Destination $imgTrain -Force
$all = Get-ChildItem -Path $imgTrain -Filter '*.jpg'
$n = $all.Count
if ($n -lt 4) { $takeVal = [Math]::Max(1,$n/4) } else { $takeVal = [Math]::Min(20,[Math]::Ceiling($n*0.1)) }
$move = $all | Select-Object -First $takeVal
foreach ($f in $move) { Move-Item -LiteralPath $f.FullName -Destination $imgVal -Force }
Write-Host ("train={0}, val={1}" -f ((Get-ChildItem $imgTrain -Filter *.jpg).Count), ((Get-ChildItem $imgVal -Filter *.jpg).Count)) -ForegroundColor Yellow

Write-Host "[4/6] 伪标注 (train/val)" -ForegroundColor Cyan
Invoke-Yolo "detect predict model=$Weights source=$imgTrain conf=0.35 save_txt=True save_conf=True project=runs/pseudo name=traffic_train device=$Device"
Invoke-Yolo "detect predict model=$Weights source=$imgVal   conf=0.35 save_txt=True save_conf=True project=runs/pseudo name=traffic_val   device=$Device"

Write-Host "[5/6] 标签映射为6类 + 修正列数" -ForegroundColor Cyan
& python scripts/remap_coco_to_six.py --src runs/pseudo/traffic_train/labels --dst $labTrain --images $imgTrain
& python scripts/remap_coco_to_six.py --src runs/pseudo/traffic_val/labels   --dst $labVal   --images $imgVal
& python scripts/fix_label_format.py $labTrain $labVal

Write-Host "[6/6] 开始训练 (20 epochs 快速验证流程)" -ForegroundColor Cyan
Invoke-Yolo "detect train model=$Weights data=datasets/yolo/zny.yaml imgsz=640 epochs=20 batch=16 device=$Device project=runs/train name=trial_from_traffic"

Write-Host "完成。最佳权重：runs/train/trial_from_traffic/weights/best.pt" -ForegroundColor Green

